# -*- coding: utf-8 -*-

"""
GPU验证脚本：检查CUDA 12.8环境配置和本地模型GPU加速可用性
包含模型预下载功能，适合国内用户使用
包含嵌入模型和大语言模型的GPU支持测试
"""

import torch
import sys
import os
import subprocess
from sentence_transformers import SentenceTransformer

# from transformers import AutoTokenizer, AutoModelForCausalLM
from modelscope import AutoModelForCausalLM, AutoTokenizer
from config import LOCAL_EM_NAME, LOCAL_EM_PATH, LOCAL_LLM_NAME, LOCAL_LLM_PATH


print("=== GPU环境验证脚本 ===")
print(f"Python版本：{sys.version}")


# 0. 提前下载模型到本地（针对国内用户优化）
def download_model_locally():
    print("\n[0] 检查并下载模型到本地...")

    # 检查本地模型是否存在
    if LOCAL_EM_PATH and os.path.exists(LOCAL_EM_PATH):
        print(f"✅ 检测到本地模型路径：{LOCAL_EM_PATH}")
        return True

    print("⚠️ 未检测到本地模型，开始下载...")
    print("   这可能需要一些时间，取决于您的网络速度")

    # 尝试使用download_models.py脚本下载
    if os.path.exists("../download_models.py"):
        print("正在使用download_models.py脚本从魔搭社区下载模型...")
        try:
            subprocess.run([sys.executable, "../download_models.py"], check=True)
            print(f"✅ 模型下载完成，已保存到：{LOCAL_EM_PATH}")
            return True
        except Exception as script_error:
            print(f"⚠️ 下载脚本执行失败：{str(script_error)}")
            print("请检查网络连接或手动运行下载脚本")

    # 如果下载脚本不存在或执行失败，不再尝试其他方式下载
    print("❌ 模型下载失败：请确保download_models.py脚本可用")
    print("\n请尝试以下解决方案：")
    print(f"1. 手动运行：python download_models.py")
    print(f"2. 或直接从魔搭社区网站下载模型：https://modelscope.cn/models")
    return False


# 1. 检查PyTorch和CUDA配置
def check_torch_cuda():
    print("\n[1] 检查PyTorch和CUDA配置...")
    print(f"PyTorch版本：{torch.__version__}")
    print(f"CUDA可用：{torch.cuda.is_available()}")

    if torch.cuda.is_available():
        torch.cuda.set_device(6)
        print(f"CUDA版本：{torch.version.cuda}")
        print(f"GPU设备数量：{torch.cuda.device_count()}")
        print(f"当前GPU：{torch.cuda.current_device()}")
        print(f"GPU名称：{torch.cuda.get_device_name(torch.cuda.current_device())}")
        print(f"GPU内存：{torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")

        # 验证CUDA版本是否为12.8
        if torch.version.cuda.startswith("12.8"):
            print("✅ 成功：CUDA版本为12.8，符合要求")
        else:
            print(f"⚠️ 警告：当前CUDA版本为{torch.version.cuda}，建议使用CUDA 12.8以获得最佳性能")
    else:
        print("❌ 错误：未检测到可用的GPU，将使用CPU运行")
        print("请确认：")
        print("1. NVIDIA驱动是否已正确安装")
        print("2. CUDA 12.8是否已正确安装")
        print("3. 是否已安装与CUDA 12.8匹配的PyTorch版本")


# 2. 检查嵌入模型是否支持GPU
def check_em_gpu_support():
    print("\n[2] 检查嵌入模型GPU支持...")
    print(f"当前配置的嵌入模型：{LOCAL_EM_NAME}")
    if LOCAL_EM_PATH and os.path.exists(LOCAL_EM_PATH):
        print(f"正在从本地路径加载嵌入模型：{LOCAL_EM_PATH}")
    else:
        print(f"⚠️ 未检测到有效本地嵌入模型路径")
        print(f"请确保嵌入模型已成功下载，或手动设置LOCAL_EM_PATH环境变量")
        print(f"当前使用的嵌入模型路径：{LOCAL_EM_PATH}")
        return False

    try:
        # 尝试加载嵌入模型
        print("正在加载嵌入模型...")

        # 加载嵌入模型，设置trust_remote_code=True以支持自定义模型
        model = SentenceTransformer(LOCAL_EM_PATH, trust_remote_code=True, local_files_only=True)
        print("✅ 模型加载成功")

        # 如果GPU可用，尝试将嵌入模型移至GPU
        if torch.cuda.is_available():
            print("尝试将嵌入模型移至GPU...")
            model = model.to("cuda")
            print("✅ 成功：嵌入模型已成功移至GPU")

            # 测试GPU上的模型推理
            print("测试GPU上的模型推理...")
            test_text = "这是一段测试文本，用于验证GPU上的模型推理功能。"
            with torch.no_grad():
                # 预热GPU
                model.encode([test_text])
                # 测量推理时间
                import time

                start_time = time.time()
                embeddings = model.encode([test_text])
                end_time = time.time()

            print(f"✅ 成功，GPU推理完成，耗时：{end_time - start_time:.4f}秒")
            print(f"生成的嵌入向量维度：{embeddings.shape[1]}")
            return True
        else:
            print("⚠️ 提示：GPU不可用，嵌入模型将在CPU上运行")
            # 即使没有GPU，仍然执行CPU上的推理测试
            test_text = "这是一段测试文本，用于验证CPU上的嵌入模型推理功能。"
            embeddings = model.encode([test_text])
            print(f"✅ CPU推理成功，生成的嵌入向量维度：{embeddings.shape[1]}")
            return False

    except Exception as e:
        print(f"❌ 错误，嵌入模型加载或测试失败：{str(e)}")
        print("可能的原因：")
        print("1. 嵌入模型文件不完整或损坏")
        print("2. 本地嵌入模型路径指向的不是有效的SentenceTransformer模型")
        print("3. 依赖库版本不兼容")
        print("\n建议解决方案：")
        print(f"1. 检查{LOCAL_EM_PATH}目录下是否有完整的嵌入模型文件")
        print(f"2. 重新运行 python download_models.py 下载模型")
        print("3. 使用最新版本的sentence_transformers和transformers库")
        return False


# 3. 检查本地大语言模型是否支持GPU
def check_llm_gpu_support():
    print("\n[3] 检查大语言模型GPU支持...")
    print(f"当前配置的大语言模型：{LOCAL_LLM_NAME}")
    if LOCAL_LLM_PATH and os.path.exists(LOCAL_LLM_PATH):
        print(f"正在从本地路径加载大语言模型：{LOCAL_LLM_PATH}")
    else:
        print(f"⚠️ 未检测到有效本地大语言模型路径")
        print(f"请确保大语言模型已成功下载，或手动设置LOCAL_LLM_PATH环境变量")
        print(f"当前使用的大语言模型路径：{LOCAL_LLM_PATH}")
        return False

    try:
        # 尝试加载大语言模型和分词器
        print("正在加载大语言模型...")

        # 首先加载分词器
        tokenizer = AutoTokenizer.from_pretrained(LOCAL_LLM_PATH, trust_remote_code=True, local_files_only=True)
        print("✅ 分词器加载成功")

        # 加载模型，设置torch_dtype为bfloat16以减少内存使用
        # model = AutoModelForCausalLM.from_pretrained(
        #     LOCAL_LLM_PATH, trust_remote_code=True, local_files_only=True, dtype=torch.bfloat16, low_cpu_mem_usage=True
        # )
        model = AutoModelForCausalLM.from_pretrained(
            LOCAL_LLM_PATH,
            dtype="auto",
            device_map="auto",
        )
        print("✅ 大语言模型加载成功")

        # 如果GPU可用，尝试将模型移至GPU
        if torch.cuda.is_available():
            print("尝试将大语言模型移至GPU...")
            model = model.to("cuda")
            print("✅ 成功：大语言模型已成功移至GPU")

            # 测试GPU上的大语言模型推理
            print("测试GPU上的大语言模型推理...")
            # test_prompt = "请严格仅根据以下指令作答，不要进行任何额外的扩展或续写。指令：请在100字内简要介绍一下人工智能。​"
            test_prompt = "请用C语言写一个冒泡排序​"
            print(f"提示：{test_prompt}")

            # 使用分词器编码输入
            inputs = tokenizer(test_prompt, return_tensors="pt").to("cuda")

            with torch.no_grad():
                # 预热GPU
                _ = model.generate(**inputs, max_new_tokens=100, do_sample=True)

                # 测量推理时间
                import time

                start_time = time.time()
                outputs = model.generate(**inputs, max_new_tokens=1000, do_sample=True, temperature=0.7, top_p=0.95)
                end_time = time.time()

            # 解码生成的文本
            generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

            print(f"✅ 成功，GPU推理完成，耗时：{end_time - start_time:.4f}秒")
            print(f"生成的文本内容：{generated_text[len(test_prompt):].strip()}")
            print("终止")
            return True
        else:
            print("⚠️ 提示：GPU不可用，大语言模型将在CPU上运行")
            print("注意：在CPU上运行大语言模型可能非常慢，建议使用GPU")
            return False

    except Exception as e:
        print(f"❌ 错误，大语言模型加载或测试失败：{str(e)}")
        print("可能的原因：")
        print("1. 大语言模型文件不完整或损坏")
        print("2. 本地大语言模型路径指向的不是有效的Hugging Face模型")
        print("3. 内存/显存不足")
        print("4. 依赖库版本不兼容")
        print("\n建议解决方案：")
        print(f"1. 检查{LOCAL_LLM_PATH}目录下是否有完整的大语言模型文件")
        print(f"2. 确保您的系统有足够的内存/显存（至少16GB内存和8GB显存）")
        print(f"3. 尝试使用量化版本的模型以减少内存使用")
        print("4. 确保安装了最新版本的transformers库")
        return False


# 4. 创建完整的验证报告
def generate_validation_report(embedding_model_supported=False, llm_model_supported=False):
    print("\n=== GPU环境验证报告 ===")
    print(f"验证时间：{__import__('datetime').datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

    if torch.cuda.is_available() and torch.version.cuda.startswith("12.8"):
        print("\n✅ 总体状态：通过")
        print("您的环境已正确配置CUDA 12.8支持，项目可以使用GPU加速。")
    elif torch.cuda.is_available():
        print("\n⚠️ 总体状态：部分通过")
        print(f"检测到GPU，但CUDA版本为{torch.version.cuda}而非12.8。")
        print("项目仍可使用GPU加速，但建议安装CUDA 12.8以获得最佳性能。")
    else:
        print("\n❌ 总体状态：未通过")
        print("未检测到可用的GPU支持。项目将在CPU上运行，性能可能较低。")
        print("请参考以下步骤进行修复：")
        print("1. 确保您的系统安装了NVIDIA GPU驱动")
        print("2. 安装CUDA 12.8和cuDNN")
        print("3. 重新安装与CUDA 12.8匹配的PyTorch版本：")
        print("   pip install torch torchvision torchaudio --index-url https://mirrors.nju.edu.cn/pytorch/whl/cu128")

    # 模型支持状态总结
    print("\n=== 模型支持状态 ===")
    if embedding_model_supported:
        print("✅ 嵌入模型：GPU支持正常")
    else:
        print("⚠️ 嵌入模型：GPU支持有问题或未测试")

    if llm_model_supported:
        print("✅ 大语言模型：GPU支持正常")
    else:
        print("⚠️ 大语言模型：GPU支持有问题或未测试")


def main():
    print(f"\n当前配置：")
    print(f"- 本地嵌入模型：{LOCAL_EM_NAME or '未设置'}")
    print(f"- 配置的本地嵌入模型路径：{LOCAL_EM_PATH or '未设置'}")
    print(f"- 本地大语言模型：{LOCAL_LLM_NAME or '未设置'}")
    print(f"- 配置的本地大语言模型路径：{LOCAL_LLM_PATH or '未设置'}")

    try:
        # 0. 下载模型到本地（针对国内用户优化）
        model_downloaded = download_model_locally()

        # 1. 检查PyTorch和CUDA配置
        check_torch_cuda()
        torch_cuda_available = torch.cuda.is_available()

        # 2. 检查嵌入模型GPU支持
        em_gpu_supported = False
        if model_downloaded:
            em_gpu_supported = check_em_gpu_support()
        else:
            print("\n[2] 检查嵌入模型GPU支持...")
            print("⚠️ 跳过GPU支持检查，因为嵌入模型下载失败")

        # 3. 检查本地大语言模型GPU支持
        llm_gpu_supported = False
        # 检查LLM路径是否存在且有效
        if LOCAL_LLM_PATH and os.path.exists(LOCAL_LLM_PATH):
            llm_gpu_supported = check_llm_gpu_support()
        else:
            print("\n[3] 检查本地大语言模型GPU支持...")
            print("⚠️ 跳过GPU支持检查，因为大语言模型路径不存在或无效")

        # 4. 生成验证报告
        generate_validation_report(em_gpu_supported, llm_gpu_supported)

        # 5. 总结
        print("\n=== 验证总结 ===")
        all_passed = torch_cuda_available and em_gpu_supported

        # 如果LLM路径存在，则要求LLM也通过测试才算全部通过
        if LOCAL_LLM_PATH and os.path.exists(LOCAL_LLM_PATH):
            all_passed = all_passed and llm_gpu_supported

        if all_passed:
            print("✅ 所有验证通过！您的环境已准备就绪。")
            return 0
        else:
            print("⚠️ 验证未完全通过，存在以下问题：")
            if not torch_cuda_available:
                print("  - PyTorch/CUDA配置有问题")
            if model_downloaded and not em_gpu_supported:
                print("  - 嵌入模型GPU支持有问题")
            if not model_downloaded:
                print("  - 嵌入模型下载失败")
            if LOCAL_LLM_PATH and os.path.exists(LOCAL_LLM_PATH) and not llm_gpu_supported:
                print("  - 大语言模型GPU支持有问题")
            print("请根据上述提示进行修复。")

            # 即使有问题也返回0，避免CI/CD流程中断
            print("\n提示：程序将以状态码0退出，方便在CI/CD环境中使用。")
            print("如果需要严格验证，请修改脚本以返回非零状态码。")
            return 0
    except Exception as e:
        print(f"\n❌ 验证过程中发生错误：{str(e)}")
        import traceback

        traceback.print_exc()
        return 0
    finally:
        print("\n=== 验证完成 ===")


if __name__ == "__main__":
    sys.exit(main())
